import os
import torch
import torchvision.transforms as transforms
from torchvision.models import resnet50, ResNet50_Weights
from PIL import Image
import faiss
import numpy as np
from tqdm import tqdm


class IndexGenerator:
    def __init__(self, use_gpu=False):
        # 设备配置
        self.device = torch.device("cuda" if use_gpu and torch.cuda.is_available() else "cpu")
        print(f"使用设备: {self.device}")

        # 加载模型（消除警告的正确方式）
        self.model = self._load_model()

        # 图像预处理
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

        # 初始化FAISS索引
        self.feature_dim = 2048
        self.index = faiss.IndexFlatL2(self.feature_dim)
        self.image_paths = []

    def _load_model(self):
        """加载ResNet-50模型（无警告版本）"""
        model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
        model = torch.nn.Sequential(*list(model.children())[:-1])  # 移除分类层
        model.to(self.device)
        model.eval()
        return model

    def extract_feature(self, image_path):
        """提取单张图片特征"""
        try:
            img = Image.open(image_path).convert("RGB")
            img_tensor = self.transform(img).unsqueeze(0).to(self.device)

            with torch.no_grad():
                feature = self.model(img_tensor)
                feature = feature.squeeze().cpu().numpy()
                return feature.reshape(1, -1)
        except Exception as e:
            print(f"跳过无效图片 {image_path}: {str(e)}")
            return None

    def build_and_save_index(self, image_dir, index_dir):
        """构建并保存索引（强制重新生成）"""
        # 确保商品图片目录存在
        if not os.path.exists(image_dir):
            raise ValueError(f"商品图片目录不存在: {image_dir}")

        # 获取所有图片文件
        image_extensions = ('.jpg', '.jpeg', '.png', '.bmp')
        image_files = []
        for root, _, files in os.walk(image_dir):
            for file in files:
                if file.lower().endswith(image_extensions):
                    image_files.append(os.path.join(root, file))

        if not image_files:
            raise ValueError(f"未在 {image_dir} 中找到任何图片，请检查目录")

        # 清空并重建索引
        self.index.reset()
        self.image_paths = []

        print(f"发现 {len(image_files)} 张图片，开始生成索引...")
        for img_path in tqdm(image_files, desc="处理图片"):
            feature = self.extract_feature(img_path)
            if feature is not None:
                self.index.add(feature)
                self.image_paths.append(img_path)

        # 保存索引
        os.makedirs(index_dir, exist_ok=True)
        faiss.write_index(self.index, os.path.join(index_dir, "image_index.index"))
        with open(os.path.join(index_dir, "image_paths.txt"), "w", encoding="utf-8") as f:
            for path in self.image_paths:
                f.write(f"{path}\n")

        print(f"索引生成完成！已保存至 {index_dir}")
        print(f"生成的文件:")
        print(f"  - image_index.index（特征索引文件）")
        print(f"  - image_paths.txt（图片路径映射）")


if __name__ == "__main__":
    # --------------------------
    # 必须修改为你的实际路径！！！
    # --------------------------
    IMAGE_DIR = "../shop/product_images"  # 商品图片目录
    INDEX_DIR = "../suoyin/image_index"  # 索引保存目录
    USE_GPU = False  # 没有NVIDIA显卡保持False

    try:
        # 强制删除旧索引（确保重新生成）
        if os.path.exists(INDEX_DIR):
            import shutil

            shutil.rmtree(INDEX_DIR)
            print(f"已删除旧索引目录: {INDEX_DIR}")

        # 生成新索引
        generator = IndexGenerator(use_gpu=USE_GPU)
        generator.build_and_save_index(IMAGE_DIR, INDEX_DIR)
    except Exception as e:
        print(f"生成索引失败: {str(e)}")
